# C3D4单元
# 公式:https://qinguoming.xyz
# -*- coding: utf-8 -*-
import numpy as np
import copy
from node import Node
from ElementBase import elementBase
from tools import *
from typing import List
from Loads import *
import scipy as sc

class C3D4(elementBase):
    """
    四面体单元,4节点线性分布,节点自由度:ux,uy,uz;节点排序和abaqus的C3d4一致;
    应力分量排序为 {11,22,33,12,23,13} ,和abaqus/explicit一致
    
    单元定义需: 节点坐标,模量E,泊松比nu,密度rho
    
    支持的单元载荷: 体积力(N/m^3);压力载荷(N/m^2);表面力载荷(N/m^2)
    
    支持集中节点力加载:F1,F2,F3
    """
    VTK_TYPE = 10
    def __init__(self, label:int,
                    node1:Node,node2:Node,node3:Node,node4:Node,
                    E:float=0.0,nu:float=0.0,rho:float=0.0):
        super().__init__(label=label,nodeDofNUM=3)
        self.node1,self.node2,self.node3,self.node4=node1,node2,node3,node4
        self.E=E
        self.nu=nu
        self.rho=rho
    @property
    def D(self):
        """返回单元各项同性弹性矩阵
        
        returns:
            ndarray(6,6), 刚度矩阵
        """
        E,nu=self.E,self.nu
        d=np.zeros((6,6))
        d[0,0]=1-nu
        d[1,1]=1-nu
        d[2,2]=1-nu
        
        d[0,1]=nu
        d[0,2]=nu
        d[1,2]=nu
        
        d[1,0]=nu
        d[2,0]=nu
        d[2,1]=nu
        
        d[3,3]=(1-2*nu)*0.5
        d[4,4]=(1-2*nu)*0.5
        d[5,5]=(1-2*nu)*0.5
        return (E/((1+nu)*(1-2*nu)))*d
    @property
    def get_connectivity(self):
        """
        获取单元的连通性编号
        
        returns:
            list[int] : 单元的节点编号
        """
        return [self.node1.label,self.node2.label,self.node3.label,self.node4.label]
    @property
    def dofnum(self)->int:
        """
        单元自由度总数
        
        returns:
            int,单元自由度数
        """
        return 12
    @property
    def NodesList(self):
        """单元节点列表
        
        returns:
            list[Node] : 节点对象列表
        """
        return [self.node1,self.node2,self.node3,self.node4]
    @property
    def ElemDofIndexs(self):
        """
        返回单元自由度列表,List[int]
        """
        return dofInds(self.node1.label,3)+dofInds(self.node2.label,3)+dofInds(self.node3.label,3)+dofInds(self.node4.label,3)
    @property
    def B(self):
        """
        计算应变矩阵B
        
        returns:
            np.ndarray (6,12), B矩阵
        """
        
        N1x,N1y,N1z=sc.linalg.solve(self.J,np.array([-1,-1,-1]))
        B1=np.array([[N1x,0,0],
                    [0,N1y,0],
                    [0,0,N1z],
                    [N1y,N1x,0],
                    [0,N1z,N1y],
                    [N1z,0,N1x]],dtype=float)
        N2x,N2y,N2z=sc.linalg.solve(self.J,np.array([1.0,0.0,0.0]))
        B2=np.array([[N2x,0,0],
                    [0,N2y,0],
                    [0,0,N2z],
                    [N2y,N2x,0],
                    [0,N2z,N2y],
                    [N2z,0,N2x]],dtype=float)
        N3x,N3y,N3z=sc.linalg.solve(self.J,np.array([0.0,1.0,0.0]))
        B3=np.array([[N3x,0,0],
                    [0,N3y,0],
                    [0,0,N3z],
                    [N3y,N3x,0],
                    [0,N3z,N3y],
                    [N3z,0,N3x]],dtype=float)
        N4x,N4y,N4z=sc.linalg.solve(self.J,np.array([0,0,1]))
        B4=np.array([[N4x,0,0],
                    [0,N4y,0],
                    [0,0,N4z],
                    [N4y,N4x,0],
                    [0,N4z,N4y],
                    [N4z,0,N4x]],dtype=float)
        B=np.concatenate((B1,B2,B3,B4),axis=1)
        if B.shape!=(6,12):
            raise ValueError("B matrix shape != (6,12)")
        return B
    @property
    def Ke(self):
        """
        计算单元刚度矩阵,C3D4单元无需高斯积分
        
        returns:
            np.ndarray(12,12) : 单元全局刚度矩阵
        """
        k=(self.detJ/6.0)*(np.dot(np.dot(self.B.T,self.D),self.B))
        return k
    @property
    def detJ(self):
        """
        计算雅可比矩阵J的行列式
        
        return:
            float : det(J)
        """
        return sc.linalg.det(self.J)
    @property
    def J(self):
        """
        计算单元雅克比矩阵J
        
        returns:
            np.ndarray(3,3) : 雅可比矩阵J (常数矩阵)
        """
        x1,y1,z1=self.node1.x,self.node1.y,self.node1.z
        x2,y2,z2=self.node2.x,self.node2.y,self.node2.z
        x3,y3,z3=self.node3.x,self.node3.y,self.node3.z
        x4,y4,z4=self.node4.x,self.node4.y,self.node4.z
        a=np.array([[-1,1,0,0],
                    [-1,0,1,0],
                    [-1,0,0,1]])
        b=np.array([[x1,y1,z1],
                    [x2,y2,z2],
                    [x3,y3,z3],
                    [x4,y4,z4]])
        return np.dot(a,b)
    @property
    def Ve(self):
        """计算单元体积
        
        return:
            float : c3d4的单元体积
        """
        return self.detJ/6.0
    @property
    def Me(self):
        """
        计算单元集中质量矩阵
        
        return:
            np.ndarray(12,12): 单元在全局坐标系下的集中质量矩阵
        """
        return np.eye(12,12)*(self.rho*self.Ve/4.0)
    def CalStress(self,ue):
        """
        计算单元的应力6个分量,c3d4单元为常应力单元
        
        args:
            ue(np.ndarray,(12,)) : 单元节点位移向量
        
        returns:
            np.ndarray(6,) : 单元应力数组
        """
        if ue.shape!=(12,1):
            ue=np.reshape(ue,(12,1))
        sigma=self.D.dot(self.CalStrain(ue))
        if sigma.shape != (6,1):
            raise ValueError("stress result shape != (6,1)")
        return  np.reshape(sigma,6)
    def CalStrain(self,ue:np.ndarray)->np.ndarray:
        """
        计算单元的应变6个分量,c3d4单元为常应变单元
        
        args:
            ue(np.ndarray,(12,)) : 单元节点位移向量
        
        returns:
            np.ndarray(6,) : 单元应变数组
        """
        if ue.shape!=(12,1):
            ue=np.reshape(ue,(12,1))
        epsilon=self.B@ue
        if epsilon.shape!=(6,1):
            raise ValueError("strain result shape != (6,1)")
        return np.reshape(epsilon,6)
    def calNormal(self,n1,n2,n3):
        """
        给定三个节点n1,n2,n3,定义向量N12,N13,计算单位法向向量n=cross(n13,n12)
        
        args:
            n1 (np.ndarray)(3,) : 节点n1的坐标数组
            n2 (np.ndarray)(3,) : 节点n2的坐标数组
            n3 (np.ndarray)(3,) : 节点n3的坐标数组
        
        returns:
            np.ndarray (3,) : 法向单位向量
        """
        # define vector
        v12,v13=n2-n1,n3-n1
        # cal the cross
        normal=np.cross(v13,v12)
        # return unit vector
        return normal/np.linalg.norm(normal)
    @property
    def Face1Normal(self):
        """
        计算单元面1(1 - 2 - 3 面)的外法向向量.normal vector=cross(vec13,vec12)
        
        returns:
            np.ndarray (3,) : 单元面1的外法向向量
        """
        return self.calNormal(self.node1.coordinate,self.node2.coordinate,self.node3.coordinate)
    @property
    def Face2Normal(self)->np.ndarray:
        """
        计算单元面2(1 - 4 - 2 面)的外法向向量.normal vector=cross(vec12,vec14)
        
        returns:
            np.ndarray (3,) : 单元面2的外法向向量
        """
        return self.calNormal(self.node1.coordinate,self.node4.coordinate,self.node2.coordinate)
    @property
    def Face3Normal(self)->np.ndarray:
        """
        计算单元面3(2 - 4 - 3 面)的外法向向量.normal vector=cross(vec23,vec24)
        
        returns:
            np.ndarray (3,) : 单元面3的外法向向量
        """
        return self.calNormal(self.node2.coordinate,self.node4.coordinate,self.node3.coordinate)
    @property
    def Face4Normal(self)->np.ndarray:
        """
        计算单元面4(3 - 4 - 1 面)的外法向向量.normal vector=cross(vec31,vec34)
        
        returns:
            np.ndarray (3,) : 单元面4的外法向向量
        """
        return self.calNormal(self.node3.coordinate,self.node4.coordinate,self.node1.coordinate)
    def ProcessELoad2EqNodeForce(self,load):
        """
        计算单元载荷的等效节点力
        
        args:
            load : 支持PressureLoad,SurfaceLoad,BodyLoad
        
        returns:
            np.ndarray(12,1) : 等效节点力向量
        """
        if isinstance(load,PressureLoad):
            return self.__ProcessELoad2EqNodeForcePressureLoad(load)
        if isinstance(load,SurfaceLoad):
            return self.__ProcessELoad2EqNodeForceSurfaceLoad(load)
        if isinstance(load,BodyLoad):
            return self.__ProcessELoad2EqNodeForceBodyLoad(load)
    def __ProcessELoad2EqNodeForceBodyLoad(self, load):
        """
        计算体力的等效节点力向量
        
        args:
            load : 支持BodyLoad
        
        returns:
            np.ndarray(12,1) : 等效节点力向量
        """
        # 体力单位N/m^3
        eqforce=np.zeros((12,1))
        p=copy.deepcopy(load)
        if p.Direction=='x':
            eqforce[::3,0]=p.LoadValue
        if p.Direction=='y':
            eqforce[1::3,0]=p.LoadValue
        if p.Direction=='z':
            eqforce[2::3,0]=p.LoadValue
        return 0.25*self.Ve*eqforce
    def __ProcessELoad2EqNodeForcePressureLoad(self,load):
        """
        计算压力载荷的等效节点力向量,单元面法向设定为外法向;
        
        args:
            load : 支持PressureLoad
        
        returns:
            np.ndarray(12,1) : 等效节点力向量
        """
        eqforce=np.zeros((12,1))
        p=copy.deepcopy(load)
        if self.label not in p.get_ElemsList:
            return eqforce
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(p.get_ElemsList, [self.label]))[0]
        # load定义中,可以一个单元的多个面都有压力,所以需要遍历
        for pos in positions:
            # 根据elem face id 进行计算
            face_id=p.FaceIdList[pos]
            if face_id==1:
                # 计算面积
                px,py,pz=-1.0*self.Face1Normal*p.PressureValue
                p0=np.array([[px],[py],[pz]])
                fs=self.TriArea(self.node1,self.node2,self.node3)
                eqforce+=(fs/3.0)*np.concatenate((p0,p0,p0,np.zeros((3,1))),axis=0)
            if face_id==2:
                px,py,pz=-1.0*self.Face2Normal*p.PressureValue
                p0=np.array([[px],[py],[pz]])
                fs=self.TriArea(self.node1,self.node4,self.node2)
                eqforce+=(fs/3.0)*np.concatenate((p0,p0,np.zeros((3,1)),p0),axis=0)
            if face_id==3:
                px,py,pz=-1.0*self.Face3Normal*p.PressureValue
                p0=np.array([[px],[py],[pz]])
                fs=self.TriArea(self.node2,self.node4,self.node3)
                eqforce+=(fs/3.0)*np.concatenate((np.zeros((3,1)),p0,p0,p0),axis=0)
            if face_id==4:
                px,py,pz=-1.0*self.Face4Normal*p.PressureValue
                p0=np.array([[px],[py],[pz]])
                fs=self.TriArea(self.node3,self.node4,self.node1)
                eqforce+=(fs/3.0)*np.concatenate((p0,np.zeros((3,1)),p0,p0),axis=0)
        if eqforce.shape!=(12,1):
            raise ValueError("eqforce.shape!=(12,1)")
        return eqforce
    def __ProcessELoad2EqNodeForceSurfaceLoad(self, load: SurfaceLoad) -> np.ndarray:
        """
        计算表面荷载(N/M^2)的等效节点荷载,作用在单元的某个面上
        
        args:
            load (SurfaceLoad) : 表面力载荷
            
        returns:
            np.ndarray(12,1) : 等效节点力向量
        """
        # 压力分布载荷单元N/m^2
        eqforce=np.zeros((12,1))
        p=copy.deepcopy(load)
        if self.label not in p.get_ElemsList:
            return eqforce
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(p.get_ElemsList, [self.label]))[0]
        # load定义中,可以一个单元的多个面都有压力,所以需要遍历
        for pos in positions:
            # 计算表面压力
            face_id=p.FaceIdList[pos]
            px,py,pz=p.DirectionVector*p.LoadValue
            p0=np.array([[px],[py],[pz]])
            if face_id==1:
                # 计算面积
                fs=self.TriArea(self.node1,self.node2,self.node3)
                eqforce+=(fs/3.0)*np.concatenate((p0,p0,p0,np.zeros((3,1))),axis=0)
            if face_id==2:
                fs=self.TriArea(self.node1,self.node4,self.node2)
                eqforce+=(fs/3.0)*np.concatenate((p0,p0,np.zeros((3,1)),p0),axis=0)
            if face_id==3:
                fs=self.TriArea(self.node2,self.node4,self.node3)
                eqforce+=(fs/3.0)*np.concatenate((np.zeros((3,1)),p0,p0,p0),axis=0)
            if face_id==4:
                fs=self.TriArea(self.node3,self.node4,self.node1)
                eqforce+=(fs/3.0)*np.concatenate((p0,np.zeros((3,1)),p0,p0),axis=0)
        return eqforce
    def TriArea(self,node1,node2,node3):
        """传入三个Node对象,计算三角形面积
        
        args:
            node1 : Node
            node2 : Node
            node3 : Node
        
        returns:
            float : 三角形面积
        """
        # 定义顶点坐标
        A = node1.coordinate
        B = node2.coordinate
        C = node3.coordinate
        # 计算向量 AB 和 AC
        AB = B - A
        AC = C - A
        # 计算叉乘
        cross_product = np.cross(AB, AC)
        # 计算叉乘的模
        cross_product_norm = np.linalg.norm(cross_product)
        # 计算面积
        return 0.5 * cross_product_norm

